

from modeller import *
from modeller.automodel import *

preys = ['prey3','prey4-1','prey5','prey13','prey15'] #this is the name of your
                                                        #input sequences in .ali format,
                                                        #which we named "preys"; all the output files will have that name.
                                                        #The order is important, since it's taken into account in the loop for the secondary structure restraints.
def temp_ali():

    log.verbose()
    env = environ()
    env.io.atom_files_directory = './:../atom_files/'

    aln = alignment(env)
    for (code, chain) in (('1EM7', 'A'), ('1FCL', 'A'), ('1FD6', 'A'), ('1GB1', 'A'), ('1PGB', 'A'), ('1QKZ', 'A'), ('2GB1', 'A'), ('2QMT', 'A'), ('2RPV', 'A'), ('3GB1', 'A')):#here are the template pdbs
        mdl = model(env, file=code, model_segment=('FIRST:'+chain, 'LAST:'+chain))
        aln.append_model(mdl, atom_files=code, align_codes=code+chain)

    for (weights, write_fit, whole) in (((1., 0., 0., 0., 1., 0.), False, True),
                                        ((1., 0.5, 1., 1., 1., 0.), False, True),
                                        ((1., 1., 1., 1., 1., 0.), True, False)):
        aln.salign(rms_cutoff=3.5, normalize_pp_scores=False,
                   rr_file='$(LIB)/as1.sim.mat', overhang=30,
                   gap_penalties_1d=(-450, -50),
                   gap_penalties_3d=(0, 3), gap_gap_score=0, gap_residue_score=0,
                   dendrogram_file='fm01645.tree',#this changes depending on the used sequence
                   alignment_type='tree', 
                                          
                                          
                   feature_weights=weights, 
                                            
                   improve_alignment=True, fit=True, write_fit=write_fit,
                   write_whole_pdb=whole, output='ALIGNMENT QUALITY')

    aln.write(file='templates.pap', alignment_format='PAP')
    aln.write(file='templates.ali', alignment_format='PIR')

    aln.salign(rms_cutoff=1.0, normalize_pp_scores=False,
               rr_file='$(LIB)/as1.sim.mat', overhang=30,
               gap_penalties_1d=(-450, -50), gap_penalties_3d=(0, 3),
               gap_gap_score=0, gap_residue_score=0, dendrogram_file='1is3A.tree',
               alignment_type='progressive', feature_weights=[0]*6,
               improve_alignment=False, fit=False, write_fit=True,
               write_whole_pdb=False, output='QUALITY')

def targ_ali():
    log.verbose()
    env = environ()

    env.libs.topology.read(file='$(LIB)/top_heav.lib')

 
    


    

    for i in preys: #this loop reads each of your .ali sequence files.
        aln = alignment(env)
        aln.append(file='templates.ali', align_codes='all')
        aln_block = len(aln)


        aln.append(file=str(i)+'.ali', align_codes=str(i))

        
        aln.salign(output='', max_gap_length=20,
                   gap_function=True,   
                   alignment_type='PAIRWISE', align_block=aln_block,
                   feature_weights=(1., 0., 0., 0., 0., 0.), overhang=0,
                   gap_penalties_1d=(-450, 0),
                   gap_penalties_2d=(0.35, 1.2, 0.9, 1.2, 0.6, 8.6, 1.2, 0., 0.),
                   similarity_flag=True)

        aln.write(file=str(i)+'-mult.ali', alignment_format='PIR')
        aln.write(file=str(i)+'mult.pap', alignment_format='PAP')

        
def first_model():
    log.verbose()
    env = environ()
    
    for i in preys: #This loop reads each multiple alignment file and takes into account the secondary structure restraints obtained elsewhere.
                    #Please note that this restraints are specific to each sequence.
        class MyModel(automodel):
            def special_restraints(self, aln):
                rsr = self.restraints

                if preys.index(i) == 0:
                    rsr.add(secondary_structure.alpha(self.residue_range('24:','44:')))
                    rsr.add(secondary_structure.alpha(self.residue_range('50:','54:')))
                    rsr.add(secondary_structure.strand(self.residue_range('9:','14:')))


                elif preys.index(i) == 1:
                    rsr.add(secondary_structure.alpha(self.residue_range('7:','7:')))
                    rsr.add(secondary_structure.alpha(self.residue_range('57:','69:')))
                    rsr.add(secondary_structure.strand(self.residue_range('12:','12:')))
                    rsr.add(secondary_structure.strand(self.residue_range('35:','41:')))
                    rsr.add(secondary_structure.strand(self.residue_range('50:','53:')))
                    rsr.add(secondary_structure.strand(self.residue_range('76:','79:')))
                    rsr.add(secondary_structure.strand(self.residue_range('85:','88:')))
                    
                elif preys.index(i) == 2:
                    rsr.add(secondary_structure.alpha(self.residue_range('25:','31:')))
                    rsr.add(secondary_structure.alpha(self.residue_range('36:','41:')))
                    rsr.add(secondary_structure.alpha(self.residue_range('57:','70:')))
                    rsr.add(secondary_structure.strand(self.residue_range('9:','14:')))
                    rsr.add(secondary_structure.strand(self.residue_range('50:','54:')))
                    rsr.add(secondary_structure.strand(self.residue_range('76:','79:')))
                    rsr.add(secondary_structure.strand(self.residue_range('85:','89:')))
                elif preys.index(i) == 3:
                    rsr.add(secondary_structure.alpha(self.residue_range('28:','41:')))
                    rsr.add(secondary_structure.alpha(self.residue_range('55:','68:')))
                    rsr.add(secondary_structure.strand(self.residue_range('8:','14:')))
                    rsr.add(secondary_structure.strand(self.residue_range('48:','51:')))
                    rsr.add(secondary_structure.strand(self.residue_range('74:','77:')))
                    rsr.add(secondary_structure.strand(self.residue_range('83:','87:')))
                elif preys.index(i) == 4:
                    rsr.add(secondary_structure.alpha(self.residue_range('26:','28:')))
                    rsr.add(secondary_structure.alpha(self.residue_range('31:','31:')))
                    rsr.add(secondary_structure.alpha(self.residue_range('33:','33:')))
                    rsr.add(secondary_structure.alpha(self.residue_range('37:','37:')))
                    rsr.add(secondary_structure.alpha(self.residue_range('39:','40:')))
                    rsr.add(secondary_structure.alpha(self.residue_range('55:','68:')))
                    rsr.add(secondary_structure.strand(self.residue_range('8:','14:')))
                    rsr.add(secondary_structure.strand(self.residue_range('48:','52:')))
                    rsr.add(secondary_structure.strand(self.residue_range('74:','77:')))
                    rsr.add(secondary_structure.strand(self.residue_range('83:','87:')))

        a = MyModel(env, alnfile=str(i)+'-mult.ali',
                      knowns=('1EM7A','1FCLA','1FD6A','1GB1A','1PGBA','1QKZA','2GB1A','2QMTA','2RPVA'), sequence=str(i))
        a.starting_model = 1
        a.ending_model = 10
        a.make()
temp_ali()
targ_ali()
first_model()
